import onnx

model_in_file = '../onnx_models/pytorch/pytorch_abs.onnx'

if __name__ == "__main__":
    model = onnx.load(model_in_file)

    nodes = model.graph.node
    nodnum = len(nodes)  # 205

    for nid in range(nodnum):
        if nodes[nid].output[0] == 'stride_32':
            print('Found stride_32: index = ', nid)
        else:
            # Print operators' name
            print(str(nid) + " : " + nodes[nid].op_type)
            # print(str(nid) + " : " + str(nodes[nid]))

    inits = model.graph.initializer
    ininum = len(inits)  # 124

    for iid in range(ininum):
        el = inits[iid]
        print('name:', el.name, ' dtype:', el.data_type, ' dim:', el.dims)
        # el.raw_data for weights and biases

    # display all the output nodes
    # print(model.graph.output)
#
# print('Done')
